feat(ptodsl): implement simt_allreduce_sum for SIMT cross-workitem all-reduce - #3
Open
kuri780 wants to merge 2 commits into
Open
feat(ptodsl): implement simt_allreduce_sum for SIMT cross-workitem all-reduce#3kuri780 wants to merge 2 commits into
kuri780 wants to merge 2 commits into
Conversation
…l-reduce Implement the pto.simt_allreduce_sum frontend interface as designed in mission/483/483_docs.md. Pure Python MLIR IR emission with three dispatch strategies: warp_reduce (<=32 threads, pow2), cross_warp_reduce (>32, pow2), ub_reduce (fallback). Supports f32 and f16. - ptodsl/ptodsl/_allreduce.py: new — 674 lines - ptodsl/ptodsl/pto.py: export simt_allreduce_sum (+3 lines) - ptodsl/tests/test_allreduce.py: new — 533 lines, all passing Co-Authored-By: Claude <noreply@anthropic.com>
Refactor _allreduce.py from helper-function outline to inline emission, add max/min reducer support, and create VPTO simulator test cases. Core changes (_allreduce.py): - Replace func.call outline with inline emission (_emit_inline) - Add simt_allreduce_max and simt_allreduce_min APIs - Add reducer dispatch tables (IDENTITY, COMBINE, REDUX) - Convert control flow to PTODSL if_() context manager - Convert ops to PTODSL wrappers (pto.const, scalar.*, redux_*, ...) - Keep raw arith only for unsigned ops (DivUIOp, RemUIOp, ShRUIOp, ult) - Auto-attach pto.simt_entry attribute for syncthreads verifier Export (pto.py): - Export simt_allreduce_sum, simt_allreduce_max, simt_allreduce_min Tests (test_allreduce.py): - Add IR structure tests for all 4 paths × 3 reducers - Add ptoas lowering verification for warp paths - Document bisheng stack-smashing bug on cross-warp scratch paths VPTO simulator tests (test/vpto/cases/micro-op/simt/allreduce_*): - 6 cases: warp_sum/max/min (32 lanes) + cross_sum/max/min (128 lanes) - kernel.pto + launch.cpp + main.cpp + golden.py + compare.py - All 6 cases verified on Ascend950PR_9599 simulator (DEVICE=SIM) Co-Authored-By: Claude <noreply@anthropic.com>
kuri780
force-pushed
the
feature-PTODSL-allreduce-clean
branch
from
June 30, 2026 08:02
761e9a1 to
a5ab39f
Compare
and0d0
pushed a commit
that referenced
this pull request
Aug 30, 2026
Resolve the deferred huge_non_headerfile finding (#3): _ops.py (5913 nbnc) is split into eight domain modules plus a re-export shim. - _ops_common.py: mode gates, normalizers, shared emitters, coercion hub - _ops_core.py: const/struct/castptr/addptr - _ops_vmem.py: vector load/store, cvt/pack contracts, gather/scatter - _ops_mask.py: predicate/mask ops and make_mask - _ops_vmath.py: vector math ops and scalar load/store - _ops_tile.py: views, allocation, and the T-op family - _ops_mte.py: MTE transfers, mad, acc-store attributes, fill_tile - _ops_simt.py: SIMT index/vote/shuffle/atomic/scalar-math/sync/pipe ops - _ops.py: re-export shim; __all__ and every module-level name (public and private, e.g. the _elements_per_vreg/vreg_type/_resolve imports used by lib/TileOps) stay importable from ptodsl._ops Code is moved verbatim by line ranges (a few helpers moved up into _ops_common to keep domain modules dependent only on the base layer; no domain-to-domain imports). test_vector_cube_ops.py patches private helpers, so its 129 patch.object(_ops, ...) sites now use a patch_ops helper that installs one mock across the whole _ops-family, preserving the pre-split patch semantics. Validated on 144 (LLVM 19.1.7): PTODSLPackage rebuilt, ctest -L PTODSL 41/41 passing, import smoke covers from ptodsl import pto, private-name imports, and TileOps.a5.thistogram.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implements
pto.simt_allreduce_sumas designed inmission/483/483_docs.md.API
pto.simt_allreduce_sum(value, *, threads, scale=1, thread_offset=0, scratch=None, scratch_offset=0) -> ScalarType
Implementation
Files
Test
python3 ptodsl/tests/test_allreduce.py
Expected: ptodsl_allreduce: PASS